
/*
 *  Checkpoint logic and helpers
 */

#include <time.h>
#include <proc.h>
#include <slab.h>
#include <string.h>
#include <sync.h>
#include <pmm.h>
#include <error.h>
#include <sched.h>
#include <elf.h>
#include <vmm.h>
#include <trap.h>
#include <unistd.h>
#include <stdio.h>
#include <sched.h>
#include <stdlib.h>
#include <assert.h>
#include <elf.h>
#include <fs.h>
#include <vfs.h>
#include <sysfile.h>
#include <file.h>
#include <swap.h>
#include <mbox.h>
#include <kio.h>
#include <sched.h>
#include <stdio.h>
#include <mp.h>
#include <checkpoint.h>
#include <clock.h>
#include <version.h>

/* unique checkpoint identifier (FIXME: should be per-container ?) */
static atomic_t ctx_count = {0}; // Hack

struct ckpt_cnt_tasks {
	struct ckpt_ctx *ctx;
	int nr;
};


/* count number of tasks in tree (and optionally fill pid's in array) */
static int __tree_count_tasks(struct proc_struct *task, void *data)
{
	struct ckpt_cnt_tasks *d = (struct ckpt_cnt_tasks *) data;
	struct ckpt_ctx *ctx = d->ctx;
	int ret;

//	ctx->tsk = task;  /* (for _ckpt_err()) */

	/* is this task cool ? */
	if (task->state == PROC_ZOMBIE)
		return 0;
	else
		ret = 1;

	if (ret < 0)
		goto out;

	if (ctx->tasks_arr) {
		if (d->nr == ctx->nr_tasks) {  /* unlikely... try again later */
			ckprintf("%(T)Bad task count (%d)\n", d->nr);
			ret = -E_BUSY;
			goto out;
		}
		ctx->tasks_arr[d->nr++] = task;
	}

	ret = 1;
 out:
//	ctx->tsk = NULL;
	return ret;
}

static int tree_count_tasks(struct ckpt_ctx *ctx)
{
	struct ckpt_cnt_tasks data;
	int ret;

	data.ctx = ctx;
	data.nr = 0;

//	ckpt_msg_lock(ctx);
	ret = walk_task_subtree(ctx->root_task, __tree_count_tasks, &data);
//	ckpt_msg_unlock(ctx);
	return ret;
}

/*
 * build_tree - scan the tasks tree in DFS order and fill in array
 * @ctx: checkpoint context
 *
 * Using DFS order simplifies the restart logic to re-create the tasks.
 *
 * On success, ctx->tasks_arr will be allocated and populated with all
 * tasks (reference taken), and ctx->nr_tasks will hold the total count.
 * The array is cleaned up by ckpt_ctx_free().
 */
static int build_tree(struct ckpt_ctx *ctx)
{
	int n, m;

	ctx->tasks_arr = NULL;

	ckprintf("Ckpt : start to build tree.\n");

	/* count tasks (no side effects) */
	n = tree_count_tasks(ctx);
	if (n < 0)
		return n;

	ckprintf("Ckpt : count tree tasks:%d\n", n);

	ctx->nr_tasks = n;
	ctx->tasks_arr = kmalloc(n * sizeof(*(ctx->tasks_arr)));
	if (!ctx->tasks_arr)
		return -E_NO_MEM;

	/* count again (now will fill array) */
	m = tree_count_tasks(ctx);

	/* unlikely, but ... (cleanup in ckpt_ctx_free) */
	if (m < 0)
		return m;
	else if (m != n)
		return -E_BUSY;

	return 0;
}

/* write the checkpoint header */
static int checkpoint_write_header(struct ckpt_ctx *ctx)
{
	struct ckpt_chunk_header chunk_h;
	struct ckpt_header h;

	init_chunk_header(&chunk_h, CKPT_CHUNK_TYPE_HEADER);
	// fill header struct
	h.arch_id = CKPT_ARCH_ID;

	// U11 : add version.h

	h.version = CHECKPOINT_VERSION;
	h.linux_version = UCORE_VERSION_CODE;

	h.uflags = ctx->uflags;
	h.time = (uint64_t)ticks;

	// get len
	chunk_h.chunk_len = sizeof(h);

	return write_chunk(ctx, &chunk_h, &h);
}

static int checkpoint_task(struct ckpt_ctx *ctx, struct proc_struct *task)
{
	write_bits(ctx, task, sizeof(struct proc_struct));
}

/* write the checkpoint pids */
//// tasks_chunk_header
//// tasks_header
//// tasks_arr
//// *tasks
static int collect_pids(struct ckpt_ctx *ctx)
{
	int i;

	// write header
	struct ckpt_chunk_header task_chunk_header;
	struct ckpt_task_header task_header;
	init_chunk_header(&task_chunk_header, CKPT_CHUNK_TYPE_TASK);
	task_chunk_header.chunk_len = sizeof(task_header);
	task_header.task_count = ctx->nr_tasks;
	write_chunk(ctx, &task_chunk_header, &task_header);

//	ckprintf("CKpt: collect pid1\n");
	// build tasks map
	map_init(&ctx->proc_map);
	for (i = 0; i < ctx->nr_tasks; i++)
		map_put(&ctx->proc_map, ctx->tasks_arr[i]);

	// write proc map table
	write_bits(ctx, &ctx->proc_map, sizeof(ctx->proc_map));

//	ckprintf("CKpt: collect pid2\n");

	// write task content
	for (i = 0; i < ctx->proc_map.len; i++)
		checkpoint_task(ctx, ctx->proc_map.arr[i].first);

//	ckprintf("CKpt: collect pid3\n");
}

// checkpoint a single mm
//// mm_dump_struct
//// mm_struct
//// ** vma_struct
//// ** vma_content
static int checkpoint_mm(struct ckpt_ctx *ctx, struct mm_struct *mm)
{
	struct ckpt_mm_dump_struct mm_dump;
	mm_dump.vma_struct_count = 0;
	list_entry_t *vma_le;

	// gather info
	for (vma_le = list_next(&(mm->mmap_list)); vma_le != &(mm->mmap_list); vma_le = list_next(vma_le))
	{
		struct vma_struct *vma = le2vma(vma_le, list_link);
		ckprintf("Cktpwr : vma : flag : %08x %08x, %08x\n", vma->vm_start, vma->vm_end, vma->vm_flags);
		mm_dump.vma_struct_count++;
	}

	// dump mm_dump_struct
	write_bits(ctx, &mm_dump, sizeof(mm_dump));
	// dump mm_struct
	write_bits(ctx, mm, sizeof(struct mm_struct));

	// switch page table for content saving
	switchtopgdir(mm);

	// dump each vma_struct
	for (vma_le = list_next(&(mm->mmap_list)); vma_le != &(mm->mmap_list); vma_le = list_next(vma_le))
	{
		struct vma_struct *vma = le2vma(vma_le, list_link);
		ckprintf("Ckpt:MM storing from %08x to %08x\n", vma->vm_start, vma->vm_end);
		write_bits(ctx, vma, sizeof(struct vma_struct));
		// dump content
		if (vma->vm_flags & VM_SHARE) // shared memory
		{
			ckprintf("Ckpt : shared mm : size %d\n", vma->shmem->len);
			ckprintf("Ckpt : shared mm And: %08x %08x %d\n", vma->shmem->shmn_cache->start, vma->shmem->shmn_cache->end, vma->shmem_off);
			if (map_get(&ctx->shared_mm_map, vma->shmem) == -1)
			{
				map_put(&ctx->shared_mm_map, vma->shmem);
				map_setvalue(&ctx->shared_mm_map, vma->shmem, vma);
				// TODO:
			}
		}
		else
			//// caution : access user memory
			write_bits(ctx, vma->vm_start, vma->vm_end - vma->vm_start);
	}

	return 0;
}

static int checkpoint_shared_memory(struct ckpt_ctx *ctx)
{
	int i;

	// write header
	struct ckpt_chunk_header sharedmm_chunk_header;
	struct ckpt_shared_mm_header sharedmm_header;
	init_chunk_header(&sharedmm_chunk_header, CKPT_CHUNK_TYPE_SHAREDMM);
	sharedmm_chunk_header.chunk_len = sizeof(sharedmm_header);
	sharedmm_header.shared_mm_count = ctx->shared_mm_map.len;
	write_chunk(ctx, &sharedmm_chunk_header, &sharedmm_header);

	/// write shm map table
	write_bits(ctx, &ctx->shared_mm_map, sizeof(ctx->shared_mm_map));

	for (i = 0; i < ctx->shared_mm_map.len; i++)
	{
		struct shmem_struct *shm = ctx->shared_mm_map.arr[i].first;
		struct vma_struct *vma = ctx->shared_mm_map.arr[i].second;
		// save shared memory
		write_bits(ctx, shm, sizeof(struct shmem_struct)); // store length
		// switch pgdir
		switchtopgdir(vma->vm_mm);
		// get address
		write_bits(ctx, vma->vm_start, vma->vm_end - vma->vm_start);

		ckprintf("Warning..........................\n");
	}
}

//
//// mm_chunk_header
//// mm_header
//// mm address array
//// **mm_checkpoint
static int collect_mm(struct ckpt_ctx *ctx)
{
	int i;

	struct mm_struct *bakmm = current->mm;

	map_init(&ctx->mm_map);
	map_init(&ctx->shared_mm_map);

	ckprintf("Ckpt: write haderr\n");

	// enum processes
	for (i = 0; i < ctx->nr_tasks; i++)
	{
		struct proc_struct *proc = ctx->tasks_arr[i];

		ckprintf("Ckpt: write ha %d\n", i);

		if (proc->mm && map_get(&ctx->mm_map, proc->mm) == -1)
			map_put(&ctx->mm_map, proc->mm);

		ckprintf("Ckpt: write ww %d\n", i);
	}

	ckprintf("Ckpt: write haderr\n");

	// write header
	struct ckpt_chunk_header mm_chunk_header;
	struct ckpt_mm_header mm_header;
	init_chunk_header(&mm_chunk_header, CKPT_CHUNK_TYPE_MM);
	mm_chunk_header.chunk_len = sizeof(mm_header);
	mm_header.mm_count = ctx->mm_map.len;
	write_chunk(ctx, &mm_chunk_header, &mm_header);

	ckprintf("Ckpt: write hader222r\n");

	// write mm map table
	write_bits(ctx, &ctx->mm_map, sizeof(struct map));

	// write mm content, collect shared memory
	for (i = 0; i < ctx->mm_map.len; i++)
	{

		ckprintf("Ckpt: write hader333r\n");
		checkpoint_mm(ctx, ctx->mm_map.arr[i].first);
	}

	ckprintf("Ckpt: write hader444r\n");

	checkpoint_shared_memory(ctx);

	switchtopgdir(bakmm);

	return 0;
}

static int checkpoint_fs(struct ckpt_ctx *ctx, struct fs_struct *fs)
{
	int i;

	// dump fs_struct
	write_bits(ctx, fs, sizeof(struct fs_struct));

//	assert(ctx->tasks_arr[0]->fs_struct == fs);
	// dump current directory
	// TODO:

	// dump filemap
	ckprintf("FS_stuc tsid se : %d %d %08x\n", FS_STRUCT_NENTRY, FD_NONE, fs);

//	write_bits(ctx, fs->filemap, sizeof(struct file) * FS_STRUCT_NENTRY);
	int sum = 0;

	// collect available filemaps
	for (i = 0; i < FS_STRUCT_NENTRY; i++)
	{
//		// dump struct
		struct file *fl = fs->filemap + i;
//
		if (i != ctx->file_d && fl->status != FD_NONE)
		{
			ckprintf("CKpt : file : %d %d %d %d\n", fl->fd, fl->status, fl->readable, fl->writable);
			sum++;
		}
	}

	write_bits(ctx, &sum, sizeof(sum));

	for (i = 0; i < FS_STRUCT_NENTRY; i++)
	{
//		// dump struct
		struct file *fl = fs->filemap + i;
//
		if (i != ctx->file_d && fl->status != FD_NONE)
			write_bits(ctx, fl, sizeof(struct file));
	}

	return 0;
}

/* write the checkpoi*/
static int collect_fs(struct ckpt_ctx *ctx)
{
	int i;

	map_init(&ctx->fs_map);

	ckprintf("Ckpt: write fs haderr\n");

	// enum processes
	for (i = 0; i < ctx->nr_tasks; i++)
	{
		struct proc_struct *proc = ctx->tasks_arr[i];

		ckprintf("Ckpt: write fs ha %d\n", i);

		if (proc->fs_struct && map_get(&ctx->fs_map, proc->fs_struct) == -1)
		{
			ckprintf("get a fs : %d %08x\n", i, proc->fs_struct);
			map_put(&ctx->fs_map, proc->fs_struct);
		}

		ckprintf("Ckpt: write ww %d\n", i);
	}

	ckprintf("Ckpt: write haderr\n");

	// write header
	struct ckpt_chunk_header fs_chunk_header;
	struct ckpt_fs_header fs_header;
	init_chunk_header(&fs_chunk_header, CKPT_CHUNK_TYPE_FS);
	fs_chunk_header.chunk_len = sizeof(fs_header);
	fs_header.fs_count = ctx->fs_map.len;
	write_chunk(ctx, &fs_chunk_header, &fs_header);

	ckprintf("Ckpt: write hader222r\n");

	// write mm map table
	write_bits(ctx, &ctx->fs_map, sizeof(struct map));

	// write mm content, collect shared memory
	for (i = 0; i < ctx->fs_map.len; i++)
	{

		ckprintf("Ckpt: write hader333r\n");
		checkpoint_fs(ctx, ctx->fs_map.arr[i].first);
	}

	ckprintf("Ckpt: write hader444r\n");

	return 0;
}

static int collect_trapframes(struct ckpt_ctx *ctx)
{
	int i;

	// write header
	struct ckpt_chunk_header tf_chunk_header;
	init_chunk_header(&tf_chunk_header, CKPT_CHUNK_TYPE_TRAPFRAME);
	tf_chunk_header.chunk_len = 0;
	write_chunk(ctx, &tf_chunk_header, &i);

	for (i = 0; i < ctx->proc_map.len; i++)
	{
		struct proc_struct *proc = (struct proc_struct*)ctx->proc_map.arr[i].first;
		write_bits(ctx, proc->tf, sizeof(struct trapframe));
	}

//	ckprintf("CKpt: collect pid3\n");
}

static int collect_timers(struct ckpt_ctx *ctx)
{
	int i;

	// write header
	struct ckpt_chunk_header timer_chunk_header;
	init_chunk_header(&timer_chunk_header, CKPT_CHUNK_TYPE_TIMER);
	timer_chunk_header.chunk_len = 0;
	write_chunk(ctx, &timer_chunk_header, &i);

	map_init(&ctx->timer_map);

	extern list_entry_t timer_list;

	list_entry_t *le = list_next(&timer_list);
	while (le != &timer_list)
	{
    	timer_t *timer = le2timer(le, timer_link);
 //   	assert(timer->expires != 0);
    	struct proc_struct *proc = timer->proc;
    	if (map_get(&ctx->proc_map, proc) != -1) // exist a proc
    	{
    		map_put(&ctx->timer_map, timer);
    	}

    	le = list_next(le);
     }

	write_bits(ctx, &ctx->timer_map, sizeof(struct map));


	for (i = 0; i < ctx->timer_map.len; i++)
	{
		ckprintf("Found a timer : .....\n");
		timer_t *timer = (timer_t*)ctx->timer_map.arr[i].first;
		write_bits(ctx, timer, sizeof(timer_t));
		// record offset

	}

//	ckprintf("CKpt: collect pid3\n");
}


/* write the checkpoint tail */
static int checkpoint_write_tail(struct ckpt_ctx *ctx)
{
return 0;
}

/* setup checkpoint-specific parts of ctx */
static int init_checkpoint_ctx(struct ckpt_ctx *ctx, pid_t pid)
{
	struct proc_struct *task;
	struct fs_struct *fs;

	/*
	 * No need for explicit cleanup here, because if an error
	 * occurs then ckpt_ctx_free() is eventually called.
	 */

	/* root task */
	ckprintf("Ckpt: init checkpoint\n");

	task = find_proc(pid);

	if (!task)
		return -E_BAD_PROC;
	else
		ctx->root_task = task;

	ckprintf("Ckpt: task struct: %08x\n", task);

//	task_lock(ctx->root_task);
	fs = ctx->root_task->fs_struct;
//	spin_lock(&fs->lock);
	ctx->root_fs_path = fs->pwd;
//	path_get(&ctx->root_fs_path);
//	spin_unlock(&fs->lock);
//	task_unlock(ctx->root_task);
	ctx->tasks_arr = NULL;
	ctx->total = 0;

	return 0;
}

// checkpoint
int do_checkpoint(struct ckpt_ctx *ctx, pid_t pid)
{
	int ret;

	ret = init_checkpoint_ctx(ctx, pid);
	if (ret < 0)
		return ret;

	ckprintf("Ckpt : init_checkpint_ctx succeed.\n");

	// count tasks and freeze them
	ret = build_tree(ctx);
	if (ret < 0)
		goto out;

	ckprintf("Ckpt: Build tree over\n");

	// write ckpt header
	ret = checkpoint_write_header(ctx);
	if (ret < 0)
		goto out;

	// checkpoint task struct
	ret = collect_pids(ctx);

	if (ret < 0)
		goto out;
	// checkpoint mm_struct
	ret = collect_mm(ctx);
	if (ret < 0)
		goto out;
	// checkpoint fs_struct
	ret = collect_fs(ctx);
	if (ret < 0)
		goto out;

	// checkpoint trapframes, including other runtime info....
	ret = collect_trapframes(ctx);
	if (ret < 0)
		goto out;

	ret = collect_timers(ctx);
	if (ret < 0)
		goto out;

	// checkpoint ...
	// write ckpt tail
	ret = checkpoint_write_tail(ctx);
	if (ret < 0)
		goto out;

	/* on success, return (unique) checkpoint identifier */
//TODO:
//	ctx->crid = atomic_inc_return(&ctx_count);
	ret = ctx->crid;
 out:
	if (ret < 0)
		ckprintf("Ckpt: failed.\n"); // report error
	else
		ckprintf("Ckpt: succeed.\n"); // succeed

	ckprintf("Ckpt total Bytes : %d\n", ctx->total);

	return ret;
}

